#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IR 运算符表达式节点定义
"""

from dataclasses import dataclass, field
from typing import List as TypingList, Optional, Iterator
import ast
from .base import IRExpr, build_expr_from_ast
from ..base import IRNode, get_default_ir_node
from ...symbol_table import TypeInfo, TypeKind


@dataclass
class IRUnaryOp(IRExpr):
    """一元运算符节点"""

    op: str = field(default="")  # 'not' or 'neg'
    operand: IRNode = field(default_factory=get_default_ir_node)
    node_type: str = "unary_op"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.operand
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.UnaryOp) -> None:
        """从 AST 一元操作节点构建 IR 一元操作节点内容

        Args:
            node: AST 一元操作节点
        """
        super().build_from_ast(node)

        # 设置操作符
        if isinstance(node.op, ast.Not):
            self.op = "not"
        elif isinstance(node.op, ast.USub):
            self.op = "neg"
        elif isinstance(node.op, ast.UAdd):
            self.op = "pos"
        elif isinstance(node.op, ast.Invert):
            self.op = "invert"
        else:
            self.op = "?"

        # 构建操作数
        self.operand = build_expr_from_ast(node.operand, self)


@dataclass
class IRBinOp(IRExpr):
    """二元运算符节点"""

    op: str = field(default="")  # '+','-','*','/','%'
    left: IRNode = field(default_factory=get_default_ir_node)
    right: IRNode = field(default_factory=get_default_ir_node)
    node_type: str = "bin_op"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.left
        yield self.right
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.BinOp) -> None:
        """从 AST 二元操作节点构建 IR 二元操作节点内容

        Args:
            node: AST 二元操作节点
        """
        super().build_from_ast(node)

        # 设置操作符
        if isinstance(node.op, ast.Add):
            self.op = "+"
        elif isinstance(node.op, ast.Sub):
            self.op = "-"
        elif isinstance(node.op, ast.Mult):
            self.op = "*"
        elif isinstance(node.op, ast.Div):
            self.op = "/"
        elif isinstance(node.op, ast.FloorDiv):
            self.op = "//"
        elif isinstance(node.op, ast.Mod):
            self.op = "%"
        elif isinstance(node.op, ast.Pow):
            self.op = "**"
        elif isinstance(node.op, ast.LShift):
            self.op = "<<"
        elif isinstance(node.op, ast.RShift):
            self.op = ">>"
        elif isinstance(node.op, ast.BitOr):
            self.op = "|"
        elif isinstance(node.op, ast.BitXor):
            self.op = "^"
        elif isinstance(node.op, ast.BitAnd):
            self.op = "&"
        elif isinstance(node.op, ast.MatMult):
            self.op = "@"
        else:
            self.op = "?"

        # 构建左操作数
        self.left = build_expr_from_ast(node.left, self)

        # 构建右操作数
        self.right = build_expr_from_ast(node.right, self)


# 为了兼容性，保留 IRBinaryOp 别名
IRBinaryOp = IRBinOp


@dataclass
class IRBoolOp(IRExpr):
    """布尔运算符节点"""

    op: str = field(default="")  # 'and' or 'or'
    values: TypingList[IRNode] = field(default_factory=list)
    node_type: str = "bool_op"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        for value in self.values:
            yield value
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.BoolOp) -> None:
        """从 AST 布尔操作节点构建 IR 布尔操作节点内容

        Args:
            node: AST 布尔操作节点
        """
        super().build_from_ast(node)

        # 设置操作符
        if isinstance(node.op, ast.And):
            self.op = "and"
        elif isinstance(node.op, ast.Or):
            self.op = "or"
        else:
            self.op = "?"

        # 构建操作数列表
        self.values = []
        for value in node.values:
            self.values.append(build_expr_from_ast(value, self))


@dataclass
class IRCompare(IRExpr):
    """比较运算符节点"""

    left: IRNode = field(default_factory=get_default_ir_node)
    ops: TypingList[str] = field(default_factory=list)  # '<','<=','>','>=','==','!='
    comparators: TypingList[IRNode] = field(default_factory=list)
    node_type: str = "compare"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.left
        for comparator in self.comparators:
            yield comparator
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.Compare) -> None:
        """从 AST 比较节点构建 IR 比较节点内容

        Args:
            node: AST 比较节点
        """
        super().build_from_ast(node)

        # 构建左操作数
        self.left = build_expr_from_ast(node.left, self)

        # 构建操作符列表
        self.ops = []
        for op in node.ops:
            if isinstance(op, ast.Eq):
                self.ops.append("==")
            elif isinstance(op, ast.NotEq):
                self.ops.append("!=")
            elif isinstance(op, ast.Lt):
                self.ops.append("<")
            elif isinstance(op, ast.LtE):
                self.ops.append("<=")
            elif isinstance(op, ast.Gt):
                self.ops.append(">")
            elif isinstance(op, ast.GtE):
                self.ops.append(">=")
            elif isinstance(op, ast.Is):
                self.ops.append("is")
            elif isinstance(op, ast.IsNot):
                self.ops.append("is not")
            elif isinstance(op, ast.In):
                self.ops.append("in")
            elif isinstance(op, ast.NotIn):
                self.ops.append("not in")
            else:
                self.ops.append("?")

        # 构建比较器列表
        self.comparators = []
        for comparator in node.comparators:
            self.comparators.append(build_expr_from_ast(comparator, self))
